"""
并行工作流代理示例 - 多源信息聚合系统
演示如何使用ParallelAgent同时执行多个独立的信息收集任务
"""

from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.tools.function_tool import FunctionTool
from google.adk.tools.tool_context import ToolContext
import json
import random
from datetime import datetime, timedelta
import os
from google.adk.models.lite_llm import LiteLlm

# API密钥配置
DEFAULT_DASHSCOPE_API_KEY = "sk-f227634bb56*************************232"  # 请替换为您的实际密钥
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", DEFAULT_DASHSCOPE_API_KEY)


# 常量定义
APP_NAME = "parallel_info_aggregator"
USER_ID = "dev_user_01"
GEMINI_MODEL = GEMINI_MODEL = LiteLlm(
                model="openai/qwen-turbo",  # 使用通义千问Turbo模型
                api_key=DASHSCOPE_API_KEY,
                api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

# 状态键定义
STATE_SEARCH_TOPIC = "search_topic"
STATE_NEWS_DATA = "news_data"
STATE_SOCIAL_DATA = "social_media_data"
STATE_MARKET_DATA = "market_data"
STATE_TECH_DATA = "tech_trends_data"
STATE_AGGREGATED_REPORT = "aggregated_report"
STATE_COLLECTION_TIMESTAMP = "collection_timestamp"

# 工具函数定义
def collect_news_data(tool_context: ToolContext):
    """收集新闻数据的工具"""
    print(f"[Tool] collect_news_data called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 防止重复调用
    call_key = f"news_data_collected_{tool_context.agent_name}"
    if session_state.get(call_key, False):
        print(f"[Debug] 新闻数据已收集，跳过重复调用")
        return session_state.get(STATE_NEWS_DATA, {})
    
    topic = session_state.get(STATE_SEARCH_TOPIC, "AI技术")
    
    # 模拟新闻数据收集
    # 在实际应用中，这里可以集成真实的新闻API，如NewsAPI, Google News等
    news_data = {
        "source": "新闻媒体",
        "topic": topic,
        "articles": [
            {
                "title": f"{topic}领域重大突破：新技术引领行业变革",
                "summary": "最新研究显示，该技术在多个应用场景中表现出色",
                "sentiment": "positive",
                "relevance_score": 0.92,
                "publish_time": "2小时前"
            },
            {
                "title": f"{topic}市场分析：投资热度持续上升",
                "summary": "分析师预测该领域将迎来新一轮增长",
                "sentiment": "positive",
                "relevance_score": 0.88,
                "publish_time": "4小时前"
            },
            {
                "title": f"{topic}应用挑战：监管政策待完善",
                "summary": "专家呼吁建立更完善的监管框架",
                "sentiment": "neutral",
                "relevance_score": 0.75,
                "publish_time": "6小时前"
            }
        ],
        "total_articles": 3,
        "average_sentiment": 0.7,
        "collection_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    session_state[STATE_NEWS_DATA] = news_data
    session_state[call_key] = True
    return news_data

def collect_social_media_data(tool_context: ToolContext):
    """收集社交媒体数据的工具"""
    print(f"[Tool] collect_social_media_data called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 防止重复调用
    call_key = f"social_data_collected_{tool_context.agent_name}"
    if session_state.get(call_key, False):
        print(f"[Debug] 社交媒体数据已收集，跳过重复调用")
        return session_state.get(STATE_SOCIAL_DATA, {})
    
    topic = session_state.get(STATE_SEARCH_TOPIC, "AI技术")
    
    # 模拟社交媒体数据收集
    # 在实际应用中，这里可以集成Twitter API, Reddit API等
    social_data = {
        "source": "社交媒体",
        "topic": topic,
        "platforms": {
            "twitter": {
                "mentions": 1250,
                "sentiment_score": 0.65,
                "trending_hashtags": [f"#{topic}", "#创新", "#技术"],
                "top_influencers": ["@tech_expert", "@ai_researcher"]
            },
            "reddit": {
                "discussions": 89,
                "upvotes": 3420,
                "comments": 567,
                "hot_topics": [f"{topic}应用", "技术讨论", "未来展望"]
            },
            "linkedin": {
                "professional_posts": 156,
                "engagement_rate": 0.78,
                "industry_insights": ["企业应用增长", "人才需求上升"]
            }
        },
        "overall_sentiment": "positive",
        "engagement_level": "high",
        "collection_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    session_state[STATE_SOCIAL_DATA] = social_data
    session_state[call_key] = True
    return social_data

def collect_market_data(tool_context: ToolContext):
    """收集市场数据的工具"""
    print(f"[Tool] collect_market_data called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 防止重复调用
    call_key = f"market_data_collected_{tool_context.agent_name}"
    if session_state.get(call_key, False):
        print(f"[Debug] 市场数据已收集，跳过重复调用")
        return session_state.get(STATE_MARKET_DATA, {})
    
    topic = session_state.get(STATE_SEARCH_TOPIC, "AI技术")
    
    # 模拟市场数据收集
    # 在实际应用中，这里可以集成股票API, 市场研究报告等
    market_data = {
        "source": "市场数据",
        "topic": topic,
        "stock_performance": {
            "related_stocks": ["NVDA", "GOOGL", "MSFT", "TSLA"],
            "average_change": "+2.3%",
            "market_cap_change": "+$15.2B",
            "trading_volume": "高于平均水平"
        },
        "investment_trends": {
            "vc_funding": "$2.1B (本季度)",
            "deal_count": 45,
            "average_deal_size": "$46.7M",
            "growth_rate": "+18% YoY"
        },
        "market_indicators": {
            "market_sentiment": "bullish",
            "volatility": "medium",
            "growth_forecast": "positive",
            "risk_level": "moderate"
        },
        "collection_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    session_state[STATE_MARKET_DATA] = market_data
    session_state[call_key] = True
    return market_data

def collect_tech_trends_data(tool_context: ToolContext):
    """收集技术趋势数据的工具"""
    print(f"[Tool] collect_tech_trends_data called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 防止重复调用
    call_key = f"tech_data_collected_{tool_context.agent_name}"
    if session_state.get(call_key, False):
        print(f"[Debug] 技术趋势数据已收集，跳过重复调用")
        return session_state.get(STATE_TECH_DATA, {})
    
    topic = session_state.get(STATE_SEARCH_TOPIC, "AI技术")
    
    # 模拟技术趋势数据收集
    # 在实际应用中，这里可以集成GitHub API, 技术论文数据库等
    tech_data = {
        "source": "技术趋势",
        "topic": topic,
        "github_activity": {
            "new_repositories": 234,
            "stars_growth": "+12.5K",
            "active_contributors": 1890,
            "popular_languages": ["Python", "JavaScript", "Go"]
        },
        "research_papers": {
            "new_publications": 67,
            "citation_growth": "+8.9%",
            "top_conferences": ["NeurIPS", "ICML", "ICLR"],
            "emerging_topics": ["多模态AI", "联邦学习", "可解释AI"]
        },
        "technology_adoption": {
            "enterprise_adoption": "75%",
            "startup_usage": "89%",
            "developer_interest": "high",
            "maturity_level": "growing"
        },
        "collection_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    session_state[STATE_TECH_DATA] = tech_data
    session_state[call_key] = True
    return tech_data

def generate_comprehensive_report(tool_context: ToolContext):
    """生成综合报告的工具"""
    print(f"[Tool] generate_comprehensive_report called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 防止重复调用
    call_key = f"report_generated_{tool_context.agent_name}"
    if session_state.get(call_key, False):
        print(f"[Debug] 综合报告已生成，跳过重复调用")
        return session_state.get(STATE_AGGREGATED_REPORT, {})
    
    # 获取所有收集的数据
    news_data = session_state.get(STATE_NEWS_DATA, {})
    social_data = session_state.get(STATE_SOCIAL_DATA, {})
    market_data = session_state.get(STATE_MARKET_DATA, {})
    tech_data = session_state.get(STATE_TECH_DATA, {})
    topic = session_state.get(STATE_SEARCH_TOPIC, "AI技术")
    
    # 生成综合报告
    report = {
        "topic": topic,
        "report_title": f"{topic}多维度信息聚合报告",
        "generation_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "data_sources": 4,
        "executive_summary": {
            "overall_sentiment": "积极",
            "market_outlook": "看涨",
            "technology_maturity": "快速发展",
            "investment_activity": "活跃"
        },
        "key_insights": [
            f"{topic}在新闻媒体中获得广泛关注，整体情绪积极",
            "社交媒体讨论热度高，专业人士参与度显著",
            "市场表现强劲，投资活动活跃",
            "技术发展迅速，开源社区贡献度高"
        ],
        "data_summary": {
            "news_articles": news_data.get("total_articles", 0),
            "social_mentions": social_data.get("platforms", {}).get("twitter", {}).get("mentions", 0),
            "market_growth": market_data.get("investment_trends", {}).get("growth_rate", "N/A"),
            "tech_repositories": tech_data.get("github_activity", {}).get("new_repositories", 0)
        },
        "recommendations": [
            "继续关注技术发展动态",
            "考虑相关投资机会",
            "加强社区参与和品牌建设",
            "密切监控市场变化"
        ]
    }
    
    session_state[STATE_AGGREGATED_REPORT] = report
    session_state[call_key] = True
    return report

def initialize_search_topic(tool_context: ToolContext):
    """初始化搜索主题的工具"""
    print(f"[Tool] initialize_search_topic called by {tool_context.agent_name}")
    
    # 正确访问session的state
    session = tool_context._invocation_context.session
    session_state = session.state
    
    # 尝试从用户输入中提取主题
    user_input = ""
    
    # 方法1: 从invocation context的user_content获取
    if hasattr(tool_context._invocation_context, 'user_content') and tool_context._invocation_context.user_content:
        user_content = tool_context._invocation_context.user_content
        if hasattr(user_content, 'parts') and user_content.parts:
            for part in user_content.parts:
                if hasattr(part, 'text') and part.text:
                    user_input = part.text.strip()
                    break
        elif user_content:
            user_input = str(user_content).strip()
    
    # 方法2: 从session的消息历史获取最新的用户消息
    if not user_input and hasattr(session, 'messages') and session.messages:
        for message in reversed(session.messages):
            if hasattr(message, 'author') and message.author == 'user':
                if hasattr(message, 'content') and message.content:
                    user_input = str(message.content).strip()
                    break
    
    print(f"[Debug] 获取到的用户输入: '{user_input}'")
    
    # 如果用户输入包含明确的主题，使用用户指定的主题
    search_topic = "AI技术"  # 默认主题
    
    if user_input:
        # 简单的主题提取逻辑
        topic_keywords = {
            "人工智能": ["人工智能", "AI", "机器学习", "深度学习"],
            "区块链": ["区块链", "比特币", "加密货币", "数字货币"],
            "量子计算": ["量子计算", "量子", "量子技术"],
            "新能源汽车": ["新能源汽车", "电动汽车", "特斯拉", "电动车"],
            "元宇宙": ["元宇宙", "虚拟现实", "VR", "AR"],
            "生物技术": ["生物技术", "基因", "医疗", "制药"]
        }
        
        user_input_lower = user_input.lower()
        for topic, keywords in topic_keywords.items():
            if any(keyword.lower() in user_input_lower for keyword in keywords):
                search_topic = topic
                break
        
        # 如果没有匹配到预定义主题，但用户输入看起来像是一个主题
        if search_topic == "AI技术" and len(user_input) < 50 and not any(char in user_input for char in "。！？"):
            search_topic = user_input
    
    # 设置搜索主题
    session_state[STATE_SEARCH_TOPIC] = search_topic
    session_state[STATE_COLLECTION_TIMESTAMP] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    print(f"[Debug] 设置搜索主题为: '{search_topic}'")
    
    return {
        "search_topic": search_topic,
        "user_input": user_input,
        "status": "initialized"
    }

# 代理定义

# 初始化代理
initialization_agent = LlmAgent(
    name="InitializationAgent",
    model=GEMINI_MODEL,
    instruction="""你是多源信息聚合系统的初始化专家。

    你的任务是：
    1. 调用initialize_search_topic工具来设置搜索主题
    2. 根据用户输入确定要分析的主题
    3. 简要说明将要进行的分析内容
    
    请先调用工具，然后告诉用户将要分析哪个主题。
    """,
    description="初始化搜索主题，准备多源数据收集",
    tools=[FunctionTool(initialize_search_topic)]
)

# 并行数据收集代理1: 新闻收集代理
news_collector_agent = LlmAgent(
    name="NewsCollectorAgent",
    model=GEMINI_MODEL,
    instruction="""你是一个专业的新闻信息收集专家。

    任务：调用collect_news_data工具收集新闻数据，然后简要总结。
    
    步骤：
    1. 调用collect_news_data工具
    2. 根据收集结果简要总结新闻趋势
    3. 完成任务
    
    只调用工具一次，然后结束。
    """,
    description="收集和分析新闻媒体中的相关信息",
    tools=[FunctionTool(collect_news_data)]
)

# 并行数据收集代理2: 社交媒体分析代理
social_media_agent = LlmAgent(
    name="SocialMediaAgent",
    model=GEMINI_MODEL,
    instruction="""你是一个社交媒体趋势分析专家。

    任务：调用collect_social_media_data工具收集社交媒体数据，然后简要分析。
    
    步骤：
    1. 调用collect_social_media_data工具
    2. 根据收集结果分析社交媒体趋势
    3. 完成任务
    
    只调用工具一次，然后结束。
    """,
    description="分析社交媒体平台上的讨论和趋势",
    tools=[FunctionTool(collect_social_media_data)]
)

# 并行数据收集代理3: 市场数据代理
market_data_agent = LlmAgent(
    name="MarketDataAgent",
    model=GEMINI_MODEL,
    instruction="""你是一个市场数据分析专家。

    任务：调用collect_market_data工具收集市场数据，然后简要分析。
    
    步骤：
    1. 调用collect_market_data工具
    2. 根据收集结果分析市场趋势
    3. 完成任务
    
    只调用工具一次，然后结束。
    """,
    description="收集和分析市场数据和投资趋势",
    tools=[FunctionTool(collect_market_data)]
)

# 并行数据收集代理4: 技术趋势代理
tech_trends_agent = LlmAgent(
    name="TechTrendsAgent",
    model=GEMINI_MODEL,
    instruction="""你是一个技术趋势分析专家。

    任务：调用collect_tech_trends_data工具收集技术数据，然后简要分析。
    
    步骤：
    1. 调用collect_tech_trends_data工具
    2. 根据收集结果分析技术趋势
    3. 完成任务
    
    只调用工具一次，然后结束。
    """,
    description="收集和分析技术发展趋势和创新动态",
    tools=[FunctionTool(collect_tech_trends_data)]
)

# 创建并行数据收集代理
parallel_data_collection = ParallelAgent(
    name="ParallelDataCollection",
    description="并行收集多源信息 - 同时执行新闻、社交媒体、市场和技术数据收集",
    sub_agents=[
        news_collector_agent,
        social_media_agent,
        market_data_agent,
        tech_trends_agent
    ]
)

# 数据聚合和报告生成代理
report_generator_agent = LlmAgent(
    name="ReportGeneratorAgent",
    model=GEMINI_MODEL,
    instruction="""你是综合信息分析和报告生成专家。

    任务：调用generate_comprehensive_report工具生成综合报告，然后提供深度分析。
    
    步骤：
    1. 调用generate_comprehensive_report工具
    2. 基于报告提供多维度洞察分析
    3. 给出趋势预测和行动建议
    4. 完成任务
    
    只调用工具一次，然后结束。
    """,
    description="整合多源数据，生成综合分析报告",
    tools=[FunctionTool(generate_comprehensive_report)]
)

# 创建完整的工作流
root_agent = SequentialAgent(
    name=APP_NAME,
    description="多源信息聚合系统 - 并行收集多维度数据并生成综合报告",
    sub_agents=[
        initialization_agent,
        parallel_data_collection,
        report_generator_agent
    ]
)

# 用于测试的辅助函数
def create_test_session_with_topic(session_service, search_topic: str):
    """创建包含搜索主题的会话"""
    session = session_service.create_session(APP_NAME, USER_ID)
    session.state[STATE_SEARCH_TOPIC] = search_topic
    session.state[STATE_COLLECTION_TIMESTAMP] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return session

# 示例搜索主题
SAMPLE_TOPICS = [
    "人工智能",
    "区块链技术",
    "量子计算",
    "新能源汽车",
    "元宇宙",
    "生物技术"
] 